#pragma once

#include <windows.h>
#include <winsock2.h>
#include <ws2spi.h>
#include <boost/thread/once.hpp>
#include <vector>

#include "XString.h"
#include "RegexDefs.h"

/*
 *=======================================================================
 * ABOUT THE HOOKS
 *
 * Introduction: Service Providers
 * -------------------------------
 * Windows Sockets 2 Architecture specifies standard sockets API for user
 * application on one hand (WSA), and API for "service providers" on the
 * other (WSP). Many providers are registered in each system, and they can be
 * identified by supported Protocol, Family and Socket Type - same parameters
 * being passed to "socket" function.
 *
 * Each provider DLL must export "WSPStartup" function, which basically
 * fills standard table of function pointers (LPWSPPROC_TABLE, "ws2spi.h")
 * with addresses of its implementation.
 *
 * In order to obtain list of TCP providers, we must:
 * 1) Invoke WSCEnumProtocols, with IPPROTO_TCP as "protocols" argument
 * 2) Find first provider with AF_INET "address family", and SOCK_STREAM type
 * 3) Use WSCGetProviderPath to get path to DLL
 *
 *
 * NOTE: There are 2 basic types of service providers: "transport" and "name
 * space". The steps mentioned above refer to the 1st type. There should be
 * an option do to similar procedure for DNS service provider, but it
 * requires additional research. As starting point you can use:
 * (-) http://developer.novell.com/support/winsock/doc/chap4pt1.htm
 * (-) MSDN: WSAEnumNameSpaceProviders [NS_DNS] or WSALookupServiceBegin
 *
 *
 *
 * How It Works
 * ------------
 * 1) We keep a list of known service prodivers' DLL, by default
 *    we only use "MSWSOCK.DLL"
 * 2) Once a known DLL is loaded, we hook its "WSPStartup" function.
 * 3) In TrapWSPStartup function, after calling the original one, we
 *    change pointers to relevant functions to our own, implemented in
 *    this file.
 *=======================================================================
 */

class WinsockHooks
{
	// list of reges'es to find service provider's DLL
	typedef std::vector<tregex> ListRegex;
	ListRegex	_service_providers_dlls;

	bool		_is_activated;
	HMODULE		_hm_service_provider;

public:
	virtual ~WinsockHooks(void) {}

	/** @return singleton instance */
	static inline WinsockHooks & GetInstance(void);

	inline void activate(void) { _is_activated = true; }

	inline void notifyDLLWasLoaded(const platform::XString &name, HMODULE hm) {
		if (_is_activated && _hm_service_provider == NULL)
			doProcessDLL(name, hm);
	}
private:
	void doProcessDLL(const platform::XString &name, HMODULE hm);

	// managing singleton instance
	WinsockHooks(void);
	static WinsockHooks * _instance;
	static boost::once_flag _instance_flag;
	static inline void CreateInstance(void) { _instance = new WinsockHooks(); }


private:
	// definitions relevant for hooking WSPStartup
	typedef int
	(WSPAPI * WSPAPI_WSPStartup)(
		WORD /*wVersionRequested*/,
		LPWSPDATA /*lpWSPData*/,
		LPWSAPROTOCOL_INFOW /*lpProtocolInfo*/,
		WSPUPCALLTABLE /*UpcallTable*/,
		LPWSPPROC_TABLE /*lpProcTable*/
    );
	WSPAPI_WSPStartup original_WSPStartup;
	WSPPROC_TABLE original_WSPFunctions;

	// the hook function
	static int WSPAPI HookWSPStartup(
		WORD wVersionRequested,
		LPWSPDATA lpWSPData,
		LPWSAPROTOCOL_INFOW lpProtocolInfo,
		WSPUPCALLTABLE UpcallTable,
		LPWSPPROC_TABLE lpProcTable
    );

	/** LPWSPCONNECT */
	static int
	WSPAPI
	TrapWSPConnect(
		SOCKET s,
		const struct sockaddr FAR * name,
		int namelen,
		LPWSABUF lpCallerData,
		LPWSABUF lpCalleeData,
		LPQOS lpSQOS,
		LPQOS lpGQOS,
		LPINT lpErrno
		);

	/** LPWSPSEND */
	static int
	WSPAPI
	TrapWSPSend(
		SOCKET s,
		LPWSABUF lpBuffers,
		DWORD dwBufferCount,
		LPDWORD lpNumberOfBytesSent,
		DWORD dwFlags,
		LPWSAOVERLAPPED lpOverlapped,
		LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
		LPWSATHREADID lpThreadId,
		LPINT lpErrno
		);

	/** LPWSPRECV */
	static int
	WSPAPI
	TrapWSPRecv(
		SOCKET s,
		LPWSABUF lpBuffers,
		DWORD dwBufferCount,
		LPDWORD lpNumberOfBytesRecvd,
		LPDWORD lpFlags,
		LPWSAOVERLAPPED lpOverlapped,
		LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
		LPWSATHREADID lpThreadId,
		LPINT lpErrno
		);

	/** LPWSPSOCKET */
	static SOCKET
	WSPAPI
	TrapWSPSocket(
		int af,
		int type,
		int protocol,
		LPWSAPROTOCOL_INFOW lpProtocolInfo,
		GROUP g,
		DWORD dwFlags,
		LPINT lpErrno
		);

	/** LPWSPLISTEN */
	static int
	WSPAPI
	TrapWSPListen(
		SOCKET s,
		int backlog,
		LPINT lpErrno
		);

	/** LPWSPSENDTO */
	static int
	WSPAPI
	TrapWSPSendTo(
		SOCKET s,
		LPWSABUF lpBuffers,
		DWORD dwBufferCount,
		LPDWORD lpNumberOfBytesSent,
		DWORD dwFlags,
		const struct sockaddr FAR * lpTo,
		int iTolen,
		LPWSAOVERLAPPED lpOverlapped,
		LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
		LPWSATHREADID lpThreadId,
		LPINT lpErrno
		);

	/** LPWSPRECVFROM */
	static int
	WSPAPI
	TrapWSPResvFrom(
		SOCKET s,
		LPWSABUF lpBuffers,
		DWORD dwBufferCount,
		LPDWORD lpNumberOfBytesRecvd,
		LPDWORD lpFlags,
		struct sockaddr FAR * lpFrom,
		LPINT lpFromlen,
		LPWSAOVERLAPPED lpOverlapped,
		LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
		LPWSATHREADID lpThreadId,
		LPINT lpErrno
		);

private:
	typedef struct hostent *
	(PASCAL * WINSOCK_gethostbyname)(const char * /*name*/);

	static struct hostent *
	PASCAL
	TrapGethostbyname(const char * name);
	
	WINSOCK_gethostbyname original_gethostbyname;
};


inline WinsockHooks & WinsockHooks::GetInstance(void) {
	boost::call_once(CreateInstance, _instance_flag);
	return *_instance;
}
